#include "rxe.h"
#include "rxe_loc.h"

struct rxe_type_info rxe_type_info[RXE_NUM_TYPES] = {
    [RXE_TYPE_UC] = {
        .name = "rxe-uc",
         //该结构体的大小，为后面申请结构体时传入参数
        .size = sizeof(struct rxe_ucontext),
         //求struct rxe_pool_entry pelem成员在struct rxe_ucontext的偏移量
        .elem_offset = offsetof(struct rxe_ucontext, pelem),
        .flags = RXE_POOL_NO_ALLOC,
    },
    [RXE_TYPE_PD] = {
        .name = "rxe-pd",
        .size = sizeof(struct rxe_pd),
        .elem_offset = offsetof(struct rxe_pd, pelem),
        .flags = RXE_POOL_NO_ALLOC,
    },
    [RXE_TYPE_AH] = {
        .name = "rxe-ah",
        .size = sizeof(struct rxe_ah),
        .elem_offset = offsetof(struct rxe_ah, pelem),
        .flags = RXE_POOL_NO_ALLOC,
    },
    [RXE_TYPE_SRQ] = {
        .name = "rxe-srq",
        .size = sizeof(struct rxe_srq),
        .elem_offset = offsetof(struct rxe_srq, pelem),
        .flags = RXE_POOL_INDEX | RXE_POOL_NO_ALLOC,
        .min_index = RXE_MIN_SRQ_INDEX,
        .max_index = RXE_MAX_SRQ_INDEX,
    },
    [RXE_TYPE_QP] = {
        .name = "rxe-qp",
        .size = sizeof(struct rxe_qp),
        .elem_offset = offsetof(struct rxe_qp, pelem),
        .flags = RXE_POOL_INDEX | RXE_POOL_NO_ALLOC,
        .min_index = RXE_MIN_SRQ_INDEX,
        .max_index = RXE_MAX_SRQ_INDEX,
    },
    [RXE_TYPE_CQ] = {
        .name = "rxe-cq",
        .size = sizeof(struct rxe_cq),
        .elem_offset = offsetof(struct rxe_cq, pelem),
        .flags = RXE_POOL_NO_ALLOC,
    },
    [RXE_TYPE_MR] = {
        .name = "rxe-mr",
        .size = sizeof(struct rxe_mr),
        .elem_offset = offsetof(struct rxe_mr, pelem),
    },
    [RXE_TYPE_MW]  = {
        .name = "rxe-mw",
        .size = sizeof(struct rxe_mw),
    },
    [RXE_TYPE_MC_GRP] = {
        .name = "rxe-mc_grp",
        .size = sizeof(struct rxe_mc_grp),
        .elem_offset = offsetof(struct rxe_mc_grp, pelem),
        .key_size = sizeof(union ib_gid),
        .flags = RXE_POOL_KEY,
        .key_offset = offsetof(struct rxe_mc_grp, mgid),
    },
    [RXE_TYPE_MC_ELEM] = {
        .name = "rxe-mc_elem",
        .size = sizeof(struct rxe_mc_elem),
        .elem_offset = offsetof(struct rxe_mc_elem, pelem),
    },
};

static inline const char *pool_name(struct rxe_pool *pool)
{
    return rxe_type_info[pool->type].name;
}

void *rxe_alloc_locked(struct rxe_pool *pool)
{
    struct rxe_type_info *info = &rxe_type_info[pool->type];
    struct rxe_pool_entry *elem;
    u8 *obj;

    if(atomic_inc_return(&pool->num_elem) > pool->max_elem)
        goto out_cnt;

    /*根据元素类型申请结构体, 如rxe_mr, rxe_pd等*/
    obj = kzalloc(info->size, GFP_ATOMIC);//GFP_ATOMIC方式申请内存
    if(!obj)
        goto out_cnt;

    //找到obj结构体里面的struct rxe_pool_entry类型的成员
    elem = (struct rxe_pool_entry *)(obj + info->elem_offset);
    //往pool里增加元素
    elem->pool = pool;
    //初始化元素的引用
    kref_init(&elem->ref_cnt);

    return obj;
out_cnt:
    //减去pool里元素的个数
    atomic_dec(&pool->num_elem);
    return NULL;
}

void *rxe_alloc(struct rxe_pool *pool)
{
    struct rxe_type_info *info = &rxe_type_info[pool->type];
    struct rxe_pool_entry *elem;
    u8 *obj;

    if(atomic_inc_return(&pool->num_elem) > pool->max_elem)
        goto out_cnt;

    /*根据元素类型申请结构体, 如rxe_mr, rxe_pd等*/
    obj = kzalloc(info->size, GFP_KERNEL);//GFP_KERNEL方式申请内存
    if(!obj)
        goto out_cnt;

    //找到obj结构体里面的struct rxe_pool_entry类型的成员
    elem = (struct rxe_pool_entry *)(obj + info->elem_offset);
    //往pool里增加元素
    elem->pool = pool;
    //初始化元素的引用--初值是1,这里的引用类似于是元素的引用，比如mr,pd等
    kref_init(&elem->ref_cnt);

    return obj;
out_cnt:
    //减去pool里元素的个数
    atomic_dec(&pool->num_elem);
    return NULL;
}

int __rxe_add_to_pool(struct rxe_pool *pool, struct rxe_pool_entry *elem)
{
    if(atomic_inc_return(&pool->num_elem) > pool->max_elem)
        goto out_cnt;

    elem->pool = pool;
    kref_init(&elem->ref_cnt);

    return 0;
out_cnt:
    atomic_dec(&pool->num_elem);
    return -EINVAL;
}

static int rxe_pool_init_index(struct rxe_pool *pool, u32 max, u32 min)
{
    int err = 0;
    size_t size;

    if((max - min + 1) < pool->max_elem) {
        pr_warn("not enough indices for max_elem\n");
        err = -EINVAL;
        goto out;
    }

    pool->index.max_index = max;
    pool->index.min_index = min;

    //https://blog.csdn.net/voice_shen/article/details/7005681
    size = BITS_TO_LONGS(max - min + 1) * sizeof(long);
    pool->index.table = kmalloc(size, GFP_KERNEL);
    if (!pool->index.table) {
        err = -ENOMEM;
        goto out;
    }

    pool->index.table_size = size;
    bitmap_zero(pool->index.table, max - min + 1);
out:
    return err;
}

int rxe_pool_init(struct rxe_dev     *rxe,
        struct rxe_pool *pool, enum rxe_elem_type type,
        unsigned int max_elem)
{
    int err = 0;
    size_t size = rxe_type_info[type].size;

    memset(pool, 0, sizeof(*pool));

    pool->rxe = rxe;
    pool->type = type;
    pool->max_elem = max_elem;
    //这个感觉没什么用？
    pool->elem_size = ALIGN(size, RXE_POOL_ALIGN);
    pool->flags = rxe_type_info[type].flags;
    //根节点初始化
    pool->index.tree = RB_ROOT;
    pool->key.tree = RB_ROOT;
    pool->cleanup = rxe_type_info[type].cleanup;

    //当前pool里面的元素个数
    atomic_set(&pool->num_elem, 0);
    rwlock_init(&pool->pool_lock);

    if(rxe_type_info[type].flags & RXE_POOL_INDEX) {
        err = rxe_pool_init_index(pool,
            rxe_type_info[type].max_index,
            rxe_type_info[type].min_index);
        if(err)
            goto out;
    }

    if(rxe_type_info[type].flags & RXE_POOL_KEY) {
        pool->key.key_offset = rxe_type_info[type].key_offset;
        pool->key.key_size = rxe_type_info[type].key_size;
    }

out:
    return err;

}

void rxe_pool_cleanup(struct rxe_pool *pool)
{
    if (atomic_read(&pool->num_elem) > 0)
        pr_warn("%s pool destroyed with unfree'd elem\n", pool_name(pool));

    kfree(pool->index.table);
}

static u32 alloc_index(struct rxe_pool *pool)
{
    u32 index;
    u32 range = pool->index.max_index - pool->index.min_index + 1;

    //https://blog.csdn.net/qq_38158479/article/details/126972203
    index = find_next_zero_bit(pool->index.table, range, pool->index.last);
    //这里如果进来，应该是超出范围了
    if(index >= range)
        index = find_first_zero_bit(pool->index.table, range);

    set_bit(index, pool->index.table);
    pool->index.last = index;
    return index + pool->index.min_index;
}

/***********************************根据index插入节点***********************************/
static int rxe_insert_index(struct rxe_pool *pool, struct rxe_pool_entry *new)
{
    struct rb_node **link = &pool->index.tree.rb_node;
    struct rb_node *parent = NULL;
    struct rxe_pool_entry *elem;

    //找到要插入节点的父节点
    while (*link) {
        parent = *link;
        elem = rb_entry(parent, struct rxe_pool_entry, index_node);

        if (elem->index == new->index) {
            pr_warn("element already exists!\n");
            return -EINVAL;
        }

        if (elem->index > new->index)
            link = &(*link)->rb_left;
        else
            link = &(*link)->rb_right;
    }
    //新节点，新节点的父节点，根节点
    rb_link_node(&new->index_node, parent, link);
    rb_insert_color(&new->index_node, &pool->index.tree);

    return 0;
}

int __rxe_add_index_locked(struct rxe_pool_entry *elem)
{
    int err;
    struct rxe_pool *pool = elem->pool;

    elem->index = alloc_index(pool);
    err = rxe_insert_index(pool, elem);

    return err;
}

int __rxe_add_index(struct rxe_pool_entry *elem)
{
    struct rxe_pool *pool = elem->pool;
    unsigned long flags;
    int err;

    write_lock_irqsave(&pool->pool_lock, flags);
    err = __rxe_add_index_locked(elem);
    write_unlock_irqrestore(&pool->pool_lock, flags);

    return err;
}

/***********************************根据index擦除节点***********************************/
void __rxe_drop_index_locked(struct rxe_pool_entry *elem)
{
    struct rxe_pool *pool = elem->pool;

    clear_bit(elem->index - pool->index.min_index, pool->index.table);
    rb_erase(&elem->index_node, &pool->index.tree);
}

void __rxe_drop_index(struct rxe_pool_entry *elem)
{
    struct rxe_pool *pool = elem->pool;
    unsigned long flags;

    write_lock_irqsave(&pool->pool_lock, flags);
    __rxe_drop_index_locked(elem);
    write_unlock_irqrestore(&pool->pool_lock, flags);
}

/***********************************根据key插入节点***********************************/
static int rxe_insert_key(struct rxe_pool *pool, struct rxe_pool_entry *new)
{
    struct rb_node **link = &pool->key.tree.rb_node;
    struct rb_node *parent = NULL;
    struct rxe_pool_entry *elem;
    int cmp;

    while (*link) {
        parent = *link;
        elem = rb_entry(parent, struct rxe_pool_entry, key_node);

        cmp = memcmp((u8 *)elem + pool->key.key_offset,
                 (u8 *)new + pool->key.key_offset, pool->key.key_size);

        if (cmp == 0) {
            pr_warn("key already exists!\n");
            return -EINVAL;
        }

        if (cmp > 0)
            link = &(*link)->rb_left;
        else
            link = &(*link)->rb_right;
    }

    rb_link_node(&new->key_node, parent, link);
    rb_insert_color(&new->key_node, &pool->key.tree);

    return 0;
}

int __rxe_add_key_locked(struct rxe_pool_entry *elem, void *key)
{
    struct rxe_pool *pool = elem->pool;
    int err;

    memcpy((u8 *)elem + pool->key.key_offset, key, pool->key.key_size);
    err = rxe_insert_key(pool, elem);

    return err;
}

int __rxe_add_key(struct rxe_pool_entry *elem, void *key)
{
    struct rxe_pool *pool = elem->pool;
    unsigned long flags;
    int err;

    write_lock_irqsave(&pool->pool_lock, flags);
    err = __rxe_add_key_locked(elem, key);
    write_unlock_irqrestore(&pool->pool_lock, flags);

    return err;
}

/***********************************根据key擦除节点***********************************/
void __rxe_drop_key_locked(struct rxe_pool_entry *elem)
{
    struct rxe_pool *pool = elem->pool;

    rb_erase(&elem->key_node, &pool->key.tree);
}

void __rxe_drop_key(struct rxe_pool_entry *elem)
{
    struct rxe_pool *pool = elem->pool;
    unsigned long flags;

    write_lock_irqsave(&pool->pool_lock, flags);
    __rxe_drop_key_locked(elem);
    write_unlock_irqrestore(&pool->pool_lock, flags);
}

/******************************************************************************/
void *rxe_pool_get_index_locked(struct rxe_pool *pool, u32 index)
{
    struct rxe_type_info *info = &rxe_type_info[pool->type];
    struct rb_node *node;
    struct rxe_pool_entry *elem;
    u8 *obj;

    node = pool->index.tree.rb_node;
    while (node) {
        elem = rb_entry(node, struct rxe_pool_entry, index_node);
        if (elem->index > index)
            node = node->rb_left;
        else if (elem->index < index)
            node = node->rb_right;
        else
            break;
    }
    //节点取出了，引用计数加一
    if (node) {
        kref_get(&elem->ref_cnt);
        obj = (u8 *)elem - info->elem_offset;
    } else {
        obj = NULL;
    }

    return obj;
}

void *rxe_pool_get_index(struct rxe_pool *pool, u32 index)
{
    u8 *obj;
    unsigned long flags;

    read_lock_irqsave(&pool->pool_lock, flags);
    obj = rxe_pool_get_index_locked(pool, index);
    read_unlock_irqrestore(&pool->pool_lock, flags);

    return obj;
}

/******************************************************************************/
void *rxe_pool_get_key_locked(struct rxe_pool *pool, void *key)
{
    struct rxe_type_info *info = &rxe_type_info[pool->type];
    struct rb_node *node;
    struct rxe_pool_entry *elem;
    u8 *obj;
    int cmp;

    node = pool->key.tree.rb_node;
    while (node) {
        elem = rb_entry(node, struct rxe_pool_entry, key_node);
        cmp = memcmp((u8 *)elem + pool->key.key_offset,
                 key, pool->key.key_size);
        if (cmp > 0)
            node = node->rb_left;
        else if (cmp < 0)
            node = node->rb_right;
        else
            break;
    }
    //节点取出了，引用计数加一
    if (node) {
        kref_get(&elem->ref_cnt);
        obj = (u8 *)elem - info->elem_offset;
    } else {
        obj = NULL;
    }
    return obj;
}

void *rxe_pool_get_key(struct rxe_pool *pool, void *key)
{
    u8 *obj;
    unsigned long flags;

    read_lock_irqsave(&pool->pool_lock, flags);
    obj = rxe_pool_get_key_locked(pool, key);
    read_unlock_irqrestore(&pool->pool_lock, flags);

    return obj;
}

void rxe_elem_release(struct kref *kref)
{
    struct rxe_pool_entry *elem =
        container_of(kref, struct rxe_pool_entry, ref_cnt);
    struct rxe_pool *pool = elem->pool;
    struct rxe_type_info *info = &rxe_type_info[pool->type];
    u8 *obj;

    if (pool->cleanup)
        pool->cleanup(elem);

    if (!(pool->flags & RXE_POOL_NO_ALLOC)) {
        obj = (u8 *)elem - info->elem_offset;
        kfree(obj);
    }
    atomic_dec(&pool->num_elem);
}
